import torch 
import torch.nn as nn 
import torch.nn.functional as F 

class ConvBNReLU(nn.Module):
    def __init__(self, nin, nout, ks=3, stride=1):
        super().__init__()
        pad = (ks-1)//2
        self.layers = nn.Sequential(
            nn.Conv2d(nin, nout, ks, stride, padding=pad), 
            nn.BatchNorm2d(nout), 
            nn.ReLU(), 
        )
    def forward(self, x):
        x = self.layers(x) 
        return x 
class DeConvBNReLU(nn.Module):
    def __init__(self, nin, nout, ks=3, stride=2):
        super().__init__()
        pad = (ks-1)//2
        self.layers = nn.Sequential(
            nn.UpsamplingNearest2d(scale_factor=stride), 
            nn.Conv2d(nin, nout, ks, 1, padding=pad), 
            nn.BatchNorm2d(nout), 
            nn.ReLU(), 
        )
    def forward(self, x):
        x = self.layers(x) 
        return x 

class Encoder(nn.Module):
    def __init__(self):
        super().__init__()
        self.layers = nn.Sequential(
            ConvBNReLU(  3,  32, 3, 2), 
            ConvBNReLU( 32,  64, 3, 2), 
            ConvBNReLU( 64, 128, 3, 2), 
            ConvBNReLU(128, 256, 3, 2), 
            DeConvBNReLU(256, 128, 3, 2), 
        )
    def forward(self, x):
        x = self.layers(x)
        return x

class Decoder(nn.Module):
    def __init__(self):
        super().__init__()
        self.layers = nn.Sequential(
            DeConvBNReLU(128, 64, 3, 2), 
            DeConvBNReLU(64, 32, 3, 2), 
            DeConvBNReLU(32, 16, 3, 2),
            nn.Conv2d(16, 3, 3, 1, padding=1), 
            nn.Sigmoid() 
        )
    def forward(self, x):
        x = self.layers(x) 
        return x 


class AutoEncoderA(nn.Module):
    def __init__(self):
        super().__init__() 
        self.encoder = Encoder()
        self.decoder = Decoder() 
    def forward(self, x):
        x = self.encoder(x) 
        x = self.decoder(x) 
        return x 
class AutoEncoderB(nn.Module):
    def __init__(self):
        super().__init__() 
        self.encoder = Encoder()
        self.decoder = Decoder() 
    def forward(self, x):
        x = self.encoder(x) 
        x = self.decoder(x) 
        return x 